"""
接受图像和直线对作为输入，其中直线对是$n\times 6$的列向量组，每一个列向量就是一个直线对，前三个元素和后三个元素构成两条直线。
"""

import torch
from torchvision.io import read_image
import numpy as np
import cv2
import os
from .calculate_base import H_solver, cross_mat


def rect_img(
    H: torch.tensor,
    path: str = None,
    src_img: torch.Tensor = None,
) -> (torch.Tensor, torch.Tensor):
    """
    将一个真实的二维物体加载至环境中
    Arguments:
      - `path` : 二维物体的路径
      - `scale` : 二维物体的缩放比例
      - `is_cached` : 是否缓存
    """
    Warning("This function is Deprecated!!!")
    if src_img:
        image = src_img
    else:
        image = read_image(path)[0, ...].to("cuda")
    # 各位置横坐标
    x = (
        (
            torch.arange(image.shape[1], dtype=torch.int16).reshape([image.shape[1], 1])
            @ torch.ones(image.shape[0], dtype=torch.int16).reshape([1, image.shape[0]])
        )
        .to("cuda")
        .flatten()
    )
    # 各位置纵坐标
    y = (
        (
            torch.ones(image.shape[0], dtype=torch.int16).reshape([image.shape[0], 1])
            @ torch.arange(image.shape[1], dtype=torch.int16).reshape(
                [1, image.shape[1]]
            )
        )
        .to("cuda")
        .flatten()
    )
    z = torch.ones_like(x, dtype=torch.float32)
    p_s = torch.stack([y, x, z], dim=1) @ H
    return p_s, image.flatten()


def back_photo(p_s: torch.Tensor, l_s: torch.Tensor) -> np.ndarray:
    """
    对相机坐标系投影，将某个给定的三维立体信息的点投影到成像平面。
    Arguments:
     - `p_s`: 一系列三维点，按列放置在一起成为一个矩阵
     - `l_s`: 一系列三维点的像素值，按列放置在一起成为一个矩阵
    """
    Warning("This function is Deprecated!!!")
    # 求得向量化的图像坐标，按行拉平
    image = torch.zeros(1200, 1200, dtype=torch.uint8, device="cuda").flatten()
    # 将像素值赋给图像
    # 这一步忽略了空间中不同位置投影到同一个像素点处的强度，默认使用最后一个点强度代替了其他所有
    p_s_proj = p_s[:, 0:2].T / torch.vstack([p_s[:, 2]] * 2)
    p_min, p_max = p_s_proj.min(), p_s_proj.max()
    p_s_proj = (p_s_proj - p_min) / (p_max - p_min) * 1200
    p_s_proj = p_s_proj.long()
    ind = (
        (p_s_proj[0, :] >= 0)
        & (p_s_proj[0, :] < 1200)
        & (p_s_proj[1, :] >= 0)
        & (p_s_proj[1, :] < 1200)
    )
    flat_ind = p_s_proj[0, ind] * 1200 + p_s_proj[1, ind]
    image[flat_ind] = l_s[ind]
    img_filtered = image.numpy(force=True).reshape([1200, 1200])
    return (
        (img_filtered - np.min(img_filtered))
        / (np.max(img_filtered) - np.min(img_filtered))
        * 255
    ).astype(np.uint8)


def rectify_main(
    lines: torch.TensorType,
    src_img: torch.Tensor = None,
    src_path: str = None,
    output_dir: str = "./output/",
    img_name: str = "rectified_img",
) -> cv2.Mat:
    """
    Arguments:
      - `lines`: 每行代表一组正交直线，前三个元素是以齐次坐标表示的第一条直线，后三个是第二条
    """
    H = (
        H_solver(
            cross_mat(torch.tensor(lines, dtype=torch.float32, device="cuda"))
            .detach()
            .clone()
        )
        .cpu()
        .numpy()
    )
    image = (
        cv2.imread(src_path, cv2.IMREAD_GRAYSCALE)
        if src_path
        else src_img.cpu().numpy()
    )
    warped_image = cv2.warpPerspective(image, H, [image.shape[1], image.shape[0]])
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    cv2.imwrite(
        os.path.join(output_dir, img_name + ".png"),
        warped_image,
    )
    return warped_image


if __name__ == "__main__":
    lines = torch.load("../cross_mat.pt")
    rectify_main(
        lines,
        src_path="../bg_img/2023-11-30-00-20-19.png",
        output_dir="../output_test/",
    )
